#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
RAG Agent - 最简短实现
"""

import json
import logging
import urllib
from typing import Dict, Any, List

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# 工具装饰器
def tool(description: str = ""):
    """工具装饰器"""
    def decorator(func):
        func._tool_description = description
        func._is_tool = True
        return func
    return decorator

class ToolResult:
    """工具执行结果"""
    def __init__(self, success: bool, data: Any, message: str):
        self.success = success
        self.data = data
        self.message = message

class RAGAgent:
    """RAG智能代理"""
    
    def __init__(self, name: str = "RAG Agent"):
        """初始化RAG Agent"""
        self.name = name
        self.tools = {}
        self._register_tools()
    
    def _register_tools(self):
        """注册所有工具"""
        for method_name in dir(self):
            method = getattr(self, method_name)
            if hasattr(method, '_is_tool') and method._is_tool:
                self.tools[method_name] = {
                    'function': method,
                    'description': getattr(method, '_tool_description', ''),
                    'name': method_name
                }
        logger.info(f"注册了 {len(self.tools)} 个工具")
    
    @tool(description="搜索知识库中的实体、关系和知识块")
    def knowledge_base_search(self, query: str, top_k: int = 5) -> ToolResult:
        """知识库搜索工具"""
        try:
            import sys, os
            sys.path.append(os.path.dirname(__file__))
            from knowledge_base_search import KnowledgeBaseSearch
            
            kb_search = KnowledgeBaseSearch()
            results = kb_search.search_all(query, top_k)
            
            formatted_results = {
                "entities": [{"title": r.title, "content": r.content} for r in results["entities"]],
                "chunks": [{"title": r.title, "content": r.content} for r in results["chunks"]]
            }
            
            return ToolResult(
                success=True,
                data=formatted_results,
                message=f"搜索完成，找到 {sum(len(v) for v in formatted_results.values())} 个结果"
            )
        except Exception as e:
            return ToolResult(success=False, data=None, message=f"搜索失败: {str(e)}")
    
    @tool(description="执行混合检索（向量检索+关键词检索）")
    def hybrid_retrieval(self, query: str, top_k: int = 5) -> ToolResult:
        """混合检索工具"""
        try:
            import sys, os
            sys.path.append(os.path.dirname(__file__))
            from hybrid_retrieval import HybridRetrieval
            
            hybrid_retrieval = HybridRetrieval()
            results = hybrid_retrieval.hybrid_search(query, top_k)
            formatted_results = [{"text": r["text"], "score": r["score"]} for r in results]
            
            return ToolResult(
                success=True,
                data=formatted_results,
                message=f"混合检索完成，找到 {len(formatted_results)} 个结果"
            )
        except Exception as e:
            return ToolResult(success=False, data=None, message=f"检索失败: {str(e)}")
    
    @tool(description="基于RAG的实时问答")
    def realtime_qa(self, question: str, top_k: int = 5) -> ToolResult:
        """实时问答工具"""
        try:
            import sys, os
            sys.path.append(os.path.dirname(__file__))
            from realtime_processor import RealLLMProcessor
            
            qa_processor = RealLLMProcessor()
            result = qa_processor.process_question(question, top_k)
            
            if result['status'] == 'success':
                return ToolResult(
                    success=True,
                    data={"answer": result['answer'], "question": result['question']},
                    message="问答完成"
                )
            else:
                return ToolResult(success=False, data=None, message=result['message'])
        except Exception as e:
            return ToolResult(success=False, data=None, message=f"问答失败: {str(e)}")
    
    @tool(description="构建和查询知识图谱")
    def knowledge_graph(self, entity_name: str, depth: int = 2) -> ToolResult:
        """知识图谱工具"""
        try:
            import sys, os
            sys.path.append(os.path.dirname(__file__))
            from knowledge_base_search import KnowledgeBaseSearch
            
            kb_search = KnowledgeBaseSearch()
            graph = kb_search.get_knowledge_graph(entity_name, depth)
            
            return ToolResult(
                success=True,
                data=graph,
                message=f"知识图谱构建完成，节点数: {len(graph['nodes'])}, 边数: {len(graph['edges'])}"
            )
        except Exception as e:
            return ToolResult(success=False, data=None, message=f"构建失败: {str(e)}")
    
    def process_user_input(self, user_input: str) -> Dict[str, Any]:
        """处理用户输入"""
        # 简单工具选择
        if "搜索" in user_input or "查找" in user_input:
            tool_name = "knowledge_base_search"
        elif "图谱" in user_input or "关系" in user_input:
            tool_name = "knowledge_graph"
        elif "检索" in user_input or "混合" in user_input:
            tool_name = "hybrid_retrieval"
        else:
            tool_name = "realtime_qa"
        
        # 执行工具
        if tool_name in self.tools:
            tool_func = self.tools[tool_name]['function']
            if tool_name == "knowledge_graph":
                result = tool_func(entity_name="人与自然", depth=2)
            else:
                result = tool_func(query=user_input, top_k=5) if tool_name != "realtime_qa" else tool_func(question=user_input, top_k=5)
            
            return {
                "content": result.data.get('answer', str(result.data)) if result.success else result.message,
                "tools_used": [tool_name],
                "confidence": 0.8 if result.success else 0.2
            }
        else:
            return {
                "content": "抱歉，无法处理您的问题",
                "tools_used": [],
                "confidence": 0.0
            }
    def get_tools_info(self) -> Dict[str, Any]:
        """获取工具信息"""
        return {
            "total_tools": len(self.tools),
            "tools": [{"name": name, "description": info["description"]} for name, info in self.tools.items()]
        }


def main():
    """测试RAG Agent"""
    print("测试RAG Agent功能")
    print("=" * 50)
    
    # 创建Agent
    agent = RAGAgent("人与自然关系专家")
    
    # 显示工具信息
    tools_info = agent.get_tools_info()
    print(f"Agent名称: {agent.name}")
    print(f"工具数量: {tools_info['total_tools']}")
    for tool in tools_info['tools']:
        print(f"  - {tool['name']}: {tool['description']}")
    
    # 测试对话
    test_questions = [
        "什么是生态危机？",
        "搜索人与自然的关系",
        "构建可持续发展的知识图谱"
    ]
    for question in test_questions:
        print(f"\n用户问题: {question}")
        response = agent.process_user_input(question)
        print(f"Agent回答: {response['content'][:100]}...")
        print(f"使用工具: {response['tools_used']}")
        print(f"置信度: {response['confidence']:.2f}")



def test_weather_api(city: str):
    """测试天气API的独立函数"""
    import urllib.request
    import urllib.parse
    
    url = 'http://apis.juhe.cn/simpleWeather/query'
    params = {
        "city": city,
        "key": "0316240f70f84bb0223785b4ff80779d",
    }
    querys = urllib.parse.urlencode(params).encode('utf-8')
    
    request = urllib.request.Request(url, data=querys)
    response = urllib.request.urlopen(request)
    content = response.read().decode('utf-8')
    
    if content:
        try:
            result = json.loads(content)
            error_code = result['error_code']
            if error_code == 0:
                temperature = result['result']['realtime']['temperature']
                humidity = result['result']['realtime']['humidity']
                info = result['result']['realtime']['info']
                wid = result['result']['realtime']['wid']
                direct = result['result']['realtime']['direct']
                power = result['result']['realtime']['power']
                aqi = result['result']['realtime']['aqi']
                print(f"温度：{temperature}°C")
                print(f"湿度：{humidity}%")
                print(f"天气：{info}")
                print(f"天气标识：{wid}")
                print(f"风向：{direct}")
                print(f"风力：{power}")
                print(f"空气质量：{aqi}")
            else:
                print(f"请求失败: {result['error_code']} {result['reason']}")
        except Exception as e:
            logger.error(str(e))
            print(f"解析结果异常：{e}")
    else:
        print("请求异常")


if __name__ == "__main__":
    print("测试天气API...")
    test_weather_api("北京")
    print("\n" + "="*50)
    main()

